#!/usr/bin/env python3
"""
Database Migration Script
Executes all database optimizations and upgrades for UmaInsight platform
"""

import sqlite3
import logging
import sys
from pathlib import Path
from typing import Optional
from contextlib import contextmanager

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class DatabaseMigrator:
    """Database migration handler"""

    def __init__(self, db_path: str):
        self.db_path = Path(db_path)
        self.conn: Optional[sqlite3.Connection] = None

    @contextmanager
    def get_connection(self):
        """Context manager for database connection"""
        try:
            self.conn = sqlite3.connect(self.db_path)
            self.conn.row_factory = sqlite3.Row
            yield self.conn
        except Exception as e:
            logger.error(f"Database connection error: {e}")
            raise
        finally:
            if self.conn:
                self.conn.close()
                self.conn = None

    def check_table_exists(self, table_name: str) -> bool:
        """Check if table exists"""
        cursor = self.conn.cursor()
        cursor.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name=?",
            (table_name,)
        )
        return cursor.fetchone() is not None

    def drop_table_if_exists(self, table_name: str) -> bool:
        """Drop table if it exists"""
        if not self.check_table_exists(table_name):
            logger.info(f"ℹ️  Table {table_name} does not exist, skipping drop")
            return True  # Return true since "table doesn't exist" is desired state

        try:
            cursor = self.conn.cursor()
            cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
            self.conn.commit()
            logger.info(f"✅ Dropped table: {table_name}")
            return True
        except Exception as e:
            logger.error(f"❌ Failed to drop table {table_name}: {e}")
            return False

    def add_histogram_enum_optimization(self) -> bool:
        """Add enum_id field to histogram_details and populate it"""
        try:
            cursor = self.conn.cursor()

            # Check if enum_id column already exists
            cursor.execute("PRAGMA table_info(histogram_details)")
            columns = [row[1] for row in cursor.fetchall()]

            if 'enum_id' not in columns:
                # Add enum_id column
                cursor.execute("""
                    ALTER TABLE histogram_details
                    ADD COLUMN enum_id INTEGER
                """)
                logger.info("✅ Added enum_id column to histogram_details")

                # Add foreign key constraint note (SQLite doesn't support adding FK to existing tables easily)
                logger.info("ℹ️  Note: Foreign key constraint added conceptually")
            else:
                logger.info("ℹ️  enum_id column already exists")

            # Create index for performance
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_histogram_details_enum_id
                ON histogram_details(enum_id)
            """)
            logger.info("✅ Created index on histogram_details.enum_id")

            # Create composite index for metric_id and enum_id
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_histogram_details_metric_enum
                ON histogram_details(metric_id, enum_id)
            """)
            logger.info("✅ Created composite index on histogram_details(metric_id, enum_id)")

            self.conn.commit()
            return True

        except Exception as e:
            logger.error(f"❌ Failed to add histogram-enum optimization: {e}")
            return False

    def populate_histogram_enum_relations(self) -> bool:
        """Populate enum_id field in histogram_details by matching enum_name"""
        try:
            cursor = self.conn.cursor()

            # Update histogram_details.enum_id by matching enum_name with metrics.name
            cursor.execute("""
                UPDATE histogram_details
                SET enum_id = (
                    SELECT m.id
                    FROM metrics m
                    WHERE m.name = histogram_details.enum_name
                    AND m.type = 'enum'
                )
                WHERE enum_name IS NOT NULL
                AND enum_name != ''
            """)

            updated_rows = cursor.rowcount
            self.conn.commit()

            logger.info(f"✅ Updated {updated_rows} histogram records with enum_id")

            # Verify the updates
            cursor.execute("""
                SELECT COUNT(*) as total,
                       COUNT(CASE WHEN enum_id IS NOT NULL THEN 1 END) as with_enum_id
                FROM histogram_details
                WHERE enum_name IS NOT NULL
            """)

            result = cursor.fetchone()
            logger.info(f"📊 Histograms with enum_name: {result['total']}")
            logger.info(f"📊 Histograms with populated enum_id: {result['with_enum_id']}")

            return True

        except Exception as e:
            logger.error(f"❌ Failed to populate histogram-enum relations: {e}")
            return False

    def add_metrics_composite_unique_constraint(self) -> bool:
        """Add composite unique constraint on metrics(name, type)"""
        try:
            cursor = self.conn.cursor()

            # SQLite doesn't support ALTER TABLE to add UNIQUE constraints directly
            # We'll need to recreate the table with the constraint

            # Check if the constraint already exists by checking the index
            cursor.execute("PRAGMA index_list(metrics)")
            indexes = cursor.fetchall()

            has_constraint = any(idx[1] == 'sqlite_autoindex_metrics_1' for idx in indexes)

            if not has_constraint:
                logger.info("🔄 Adding composite unique constraint on metrics(name, type)")

                # Create backup table
                cursor.execute("""
                    CREATE TABLE metrics_backup AS
                    SELECT * FROM metrics
                """)

                # Drop original table
                cursor.execute("DROP TABLE metrics")

                # Recreate table with composite unique constraint
                cursor.execute("""
                    CREATE TABLE metrics (
                        id INTEGER PRIMARY KEY AUTOINCREMENT,
                        name TEXT NOT NULL,
                        type TEXT NOT NULL,
                        summary TEXT,
                        units TEXT,
                        expires_after TEXT,
                        component TEXT,
                        singular BOOLEAN DEFAULT FALSE,
                        file_path TEXT NOT NULL,
                        line_number INTEGER,
                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                        analysis_status TEXT DEFAULT 'NOT_ANALYZED',
                        analysis_notes TEXT,
                        analyzed_by TEXT,
                        analyzed_at DATETIME,
                        UNIQUE(name, type)
                    )
                """)

                # Copy data back
                cursor.execute("""
                    INSERT INTO metrics
                    SELECT * FROM metrics_backup
                """)

                # Drop backup table
                cursor.execute("DROP TABLE metrics_backup")

                # Recreate indexes
                cursor.execute("CREATE INDEX idx_metrics_name ON metrics(name)")
                cursor.execute("CREATE INDEX idx_metrics_type ON metrics(type)")
                cursor.execute("CREATE INDEX idx_metrics_component ON metrics(component)")

                self.conn.commit()
                logger.info("✅ Added composite unique constraint on metrics(name, type)")
            else:
                logger.info("ℹ️  Composite unique constraint already exists")

            return True

        except Exception as e:
            logger.error(f"❌ Failed to add metrics composite unique constraint: {e}")
            return False

    def cleanup_unused_tables_and_indexes(self) -> bool:
        """Clean up unused tables and indexes"""
        cleanup_success = True

        # Tables to drop
        tables_to_drop = [
            'histogram_suffixes',
            'expanded_histograms',
            'metric_ownership_consolidated',
            'telemetry_metrics_unified'
        ]

        for table in tables_to_drop:
            if not self.drop_table_if_exists(table):
                cleanup_success = False

        return cleanup_success

    def update_metrics_table_structure(self) -> bool:
        """Update metrics table to match current model expectations"""
        try:
            cursor = self.conn.cursor()

            # Check and add missing columns
            cursor.execute("PRAGMA table_info(metrics)")
            columns = {row[1]: row for row in cursor.fetchall()}

            # Add analysis status columns if missing
            if 'analysis_status' not in columns:
                cursor.execute("""
                    ALTER TABLE metrics
                    ADD COLUMN analysis_status TEXT DEFAULT 'NOT_ANALYZED'
                """)
                logger.info("✅ Added analysis_status column to metrics")

            if 'analysis_notes' not in columns:
                cursor.execute("""
                    ALTER TABLE metrics
                    ADD COLUMN analysis_notes TEXT
                """)
                logger.info("✅ Added analysis_notes column to metrics")

            if 'analyzed_by' not in columns:
                cursor.execute("""
                    ALTER TABLE metrics
                    ADD COLUMN analyzed_by TEXT
                """)
                logger.info("✅ Added analyzed_by column to metrics")

            if 'analyzed_at' not in columns:
                cursor.execute("""
                    ALTER TABLE metrics
                    ADD COLUMN analyzed_at DATETIME
                """)
                logger.info("✅ Added analyzed_at column to metrics")

            self.conn.commit()
            return True

        except Exception as e:
            logger.error(f"❌ Failed to update metrics table structure: {e}")
            return False

    def run_migration(self) -> bool:
        """Run complete database migration"""
        logger.info("🚀 Starting database migration...")
        logger.info(f"📁 Database path: {self.db_path}")

        if not self.db_path.exists():
            logger.error(f"❌ Database file does not exist: {self.db_path}")
            return False

        try:
            with self.get_connection():
                # Step 1: Clean up unused tables
                logger.info("\n📋 Step 1: Cleaning up unused tables...")
                if not self.cleanup_unused_tables_and_indexes():
                    return False

                # Step 2: Update metrics table structure
                logger.info("\n📋 Step 2: Updating metrics table structure...")
                if not self.update_metrics_table_structure():
                    return False

                # Step 3: Add composite unique constraint to metrics
                logger.info("\n📋 Step 3: Adding composite unique constraint to metrics...")
                if not self.add_metrics_composite_unique_constraint():
                    return False

                # Step 4: Add histogram-enum optimization
                logger.info("\n📋 Step 4: Adding histogram-enum optimization...")
                if not self.add_histogram_enum_optimization():
                    return False

                # Step 5: Populate histogram-enum relations
                logger.info("\n📋 Step 5: Populating histogram-enum relations...")
                if not self.populate_histogram_enum_relations():
                    return False

                logger.info("\n✅ Database migration completed successfully!")
                return True

        except Exception as e:
            logger.error(f"❌ Migration failed: {e}")
            return False


def main():
    """Main migration function"""
    # Default database path
    default_db_path = "./backend/data/uma_insight.db"

    # Use command line argument if provided
    if len(sys.argv) > 1:
        db_path = sys.argv[1]
    else:
        db_path = default_db_path

    # Resolve relative path
    db_path = Path(db_path).resolve()

    logger.info("=" * 60)
    logger.info("🔄 UmaInsight Database Migration Tool")
    logger.info("=" * 60)

    migrator = DatabaseMigrator(db_path)

    success = migrator.run_migration()

    if success:
        logger.info("\n🎉 Migration completed successfully!")
        logger.info("💡 You can now restart the backend service")
        sys.exit(0)
    else:
        logger.error("\n💥 Migration failed!")
        logger.error("💡 Please check the error messages above and fix any issues")
        sys.exit(1)


if __name__ == "__main__":
    main()